package com.spx.chapter09;

import com.spx.chapter05.pojo.Event;
import com.spx.util.ClickSource;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.AggregatingState;
import org.apache.flink.api.common.state.AggregatingStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

import java.time.Duration;

/**
 * create by undeRdoG on  2022-05-06  22:01
 * 凡心所向，素履以往，生如逆旅，一苇以航。
 */
public class AggregateTimeStampExample {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        SingleOutputStreamOperator<Event> stream = env.addSource(new ClickSource())
                .assignTimestampsAndWatermarks(
                        WatermarkStrategy.<Event>forBoundedOutOfOrderness(Duration.ZERO)
                                .withTimestampAssigner(new SerializableTimestampAssigner<Event>() {
                                    @Override
                                    public long extractTimestamp(Event element, long recordTimestamp) {
                                        return element.timestamp;
                                    }
                                })
                );


        // 自定义实现平均时间戳

        stream.print("Stream");


        stream.keyBy(data -> data.user)
                .flatMap(new AvgTimeFunction())
                .print("res");

        env.execute();
    }


    public static class AvgTimeFunction extends RichFlatMapFunction<Event, String> {


        // 定义一个聚合的状态，用来保存平均时间戳
        private AggregatingState<Event, Long> aggState;

        // 定义一个常量保存个数
        private ValueState<Long> countState;


        @Override
        public void open(Configuration parameters) throws Exception {
            this.countState = getRuntimeContext().getState(new ValueStateDescriptor<Long>("countState", Long.class));
            this.aggState = getRuntimeContext().getAggregatingState(new AggregatingStateDescriptor<Event, Tuple2<Long, Long>, Long>("aggState", new AggregateFunction<Event, Tuple2<Long, Long>, Long>() {
                @Override
                public Tuple2<Long, Long> createAccumulator() {
                    return Tuple2.of(0L, 0L);
                }

                @Override
                public Tuple2<Long, Long> add(Event value, Tuple2<Long, Long> accumulator) {
                    accumulator.setFields(accumulator.f0 + value.timestamp, accumulator.f1 + 1L);
                    return accumulator;
                }

                @Override
                public Long getResult(Tuple2<Long, Long> accumulator) {
                    return accumulator.f0 / accumulator.f1;
                }

                @Override
                public Tuple2<Long, Long> merge(Tuple2<Long, Long> a, Tuple2<Long, Long> b) {
                    return null;
                }
            }, Types.TUPLE(Types.LONG, Types.LONG)));
        }

        @Override
        public void flatMap(Event value, Collector<String> out) throws Exception {

            if (this.countState.value() == null){
                this.countState.update(1L);
                this.aggState.add(value);
            }

            if (this.countState.value().equals(5L)){
                out.collect(this.aggState.get().toString());
                this.countState.clear();
                this.aggState.clear();
            }else {
                Long currentCount = this.countState.value();
                this.countState.update(currentCount + 1L);
                this.aggState.add(value);
            }

        }
    }
}
